package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/glycerine/rbtree"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"
	"net"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unsafe"
)

type ServerConfiguration struct {
	ServerPort          uint16                         `yaml:"server-port"`
	HostReadTimeout     string                         `yaml:"host-read-timeout"`
	HostWriteTimeout    string                         `yaml:"host-write-timeout"`
	ProxyReadTimeout    string                         `yaml:"proxy-read-timeout"`
	ProxyWriteTimeout   string                         `yaml:"proxy-write-timeout"`
	SignTimeout         string                         `yaml:"sign-timeout"`
	MaxPacketBodySize   int                            `yaml:"max-packet-body-size"`
	TransportBufferSize int                            `yaml:"transport-buffer-size"`
	Hosts               map[string]*HostConfiguration  `yaml:"hosts"`
	Proxies             map[uint16]*ProxyConfiguration `yaml:"proxies"`
}

type HostConfiguration struct {
	SecretKey    string `yaml:"secret-key"`
	ReadTimeout  string `yaml:"read-timeout"`
	WriteTimeout string `yaml:"write-timeout"`
}

type ProxyConfiguration struct {
	Host          string `yaml:"host"`
	Network       string `yaml:"network"`
	TargetAddress string `yaml:"target-address"`
	ReadTimeout   string `yaml:"read-timeout"`
	WriteTimeout  string `yaml:"write-timeout"`
}

type Server struct {
	ServerPort          uint16
	HostReadTimeout     time.Duration
	HostWriteTimeout    time.Duration
	ProxyReadTimeout    time.Duration
	ProxyWriteTimeout   time.Duration
	SignTimeout         time.Duration
	MaxPacketBodySize   int
	TransportBufferSize int
	Listener            net.Listener
	Hosts               map[string]*Host
	Proxies             map[uint16]*Proxy
}

type Host struct {
	Server       *Server
	HostName     string
	SecretKey    string
	ReadTimeout  time.Duration
	WriteTimeout time.Duration
	LineSequence int32
	Lines        *rbtree.Tree
	LineTreeLock *sync.Mutex
	Connected    int32
	Conn         net.Conn
}

type Proxy struct {
	Server        *Server
	ProxyPort     uint16
	Host          *Host
	Network       string
	TargetAddress string
	ReadTimeout   time.Duration
	WriteTimeout  time.Duration
	Listener      net.Listener
}

type Line struct {
	Server      *Server
	Proxy       *Proxy
	Host        *Host
	LineNum     int32
	ProxyConn   net.Conn
	HostConn    net.Conn
	Connected   int32
	ConnChannel chan bool
	Err         interface{}
}

type ServerConnContext struct {
	Server        *Server
	Host          *Host
	HostConnected bool
	Conn          net.Conn
	Err           interface{}
	PktHeader     *PacketHeader
}

type ConnHandler func(obj interface{}, conn net.Conn)

var (
	kIsShowHelp   bool
	kConfigPath   string
	kServerConfig *ServerConfiguration
	kServer       *Server
)

func init() {
	flag.StringVar(&kConfigPath, "c", "rainbow-server.yml", "Set the configuration `filename`")
	flag.BoolVar(&kIsShowHelp, "h", false, "Show help")
	flag.Usage = usage
}

func initConfig() {
	var err error

	flag.Parse()
	if kIsShowHelp {
		usage()
		os.Exit(0)
	}

	kServerConfig = &ServerConfiguration{
		ServerPort:          9900,
		Hosts:               map[string]*HostConfiguration{},
		Proxies:             map[uint16]*ProxyConfiguration{},
		HostReadTimeout:     "3m",
		HostWriteTimeout:    "3m",
		ProxyReadTimeout:    "3m",
		ProxyWriteTimeout:   "3m",
		SignTimeout:         "30s",
		MaxPacketBodySize:   65536 - PacketHeaderSize,
		TransportBufferSize: 65536,
	}

	file, err := os.Open(kConfigPath)
	if err != nil {
		panic(err)
	}
	defer closeFile(file)

	err = yaml.NewDecoder(file).Decode(&kServerConfig)
	if err != nil {
		panic(err)
	}

	configContent, err := yaml.Marshal(kServerConfig)
	if err != nil {
		panic(err)
	}

	fmt.Println(string(configContent))

	kServer = &Server{
		ServerPort:          kServerConfig.ServerPort,
		HostReadTimeout:     parseDurationConfig("host-read-timeout", kServerConfig.HostReadTimeout),
		HostWriteTimeout:    parseDurationConfig("host-write-timeout", kServerConfig.HostWriteTimeout),
		ProxyReadTimeout:    parseDurationConfig("proxy-read-timeout", kServerConfig.ProxyReadTimeout),
		ProxyWriteTimeout:   parseDurationConfig("proxy-write-timeout", kServerConfig.ProxyWriteTimeout),
		SignTimeout:         parseDurationConfig("sign-timeout", kServerConfig.SignTimeout),
		MaxPacketBodySize:   kServerConfig.MaxPacketBodySize,
		TransportBufferSize: kServerConfig.TransportBufferSize,
		Hosts:               map[string]*Host{},
		Proxies:             map[uint16]*Proxy{},
	}

	for hostName := range kServerConfig.Hosts {
		objName := fmt.Sprintf("%s.%s", "clients", hostName)
		hostConfig := kServerConfig.Hosts[hostName]
		secretKey := hostConfig.SecretKey
		readTimeout := hostConfig.ReadTimeout
		writeTimeout := hostConfig.WriteTimeout
		host := &Host{
			Server:    kServer,
			HostName:  hostName,
			SecretKey: secretKey,
			ReadTimeout: parseObjDurationConfig(
				objName, "read-timeout", readTimeout, kServer.HostReadTimeout),
			WriteTimeout: parseObjDurationConfig(
				objName, "write-timeout", writeTimeout, kServer.HostWriteTimeout),
			LineSequence: 0,
			Lines:        rbtree.NewTree(lineCompare),
			LineTreeLock: &sync.Mutex{},
		}
		kServer.Hosts[hostName] = host
	}

	for proxyPort := range kServerConfig.Proxies {
		objName := fmt.Sprintf("%s.%d", "proxies", proxyPort)
		proxyConfig := kServerConfig.Proxies[proxyPort]
		readTimeout := proxyConfig.ReadTimeout
		writeTimeout := proxyConfig.WriteTimeout
		host := kServer.Hosts[proxyConfig.Host]
		if host == nil {
			panic(errors.New(fmt.Sprintf(
				"Could not found host '%s' for proxy '%d'", proxyConfig.Host, proxyPort)))
		}
		proxy := &Proxy{
			Server:        kServer,
			ProxyPort:     proxyPort,
			Host:          host,
			Network:       proxyConfig.Network,
			TargetAddress: proxyConfig.TargetAddress,
			ReadTimeout: parseObjDurationConfig(
				objName, "read-timeout", readTimeout, kServer.ProxyReadTimeout),
			WriteTimeout: parseObjDurationConfig(
				objName, "write-timeout", writeTimeout, kServer.ProxyWriteTimeout),
		}
		kServer.Proxies[proxyPort] = proxy
	}
}

func usage() {
	_, _ = fmt.Fprintln(os.Stderr, "rainbow-server version: 1.0")
	_, _ = fmt.Fprintln(os.Stderr, "Usage: rainbow-server [-c filename]")
	_, _ = fmt.Fprintln(os.Stderr, "Options: ")
	flag.PrintDefaults()
}

func listen() {
	kServer.Listener = listenPort(kServer.ServerPort)
	for proxyPort := range kServer.Proxies {
		proxy := kServer.Proxies[proxyPort]
		proxy.Listener = listenPort(proxyPort)
		go acceptConn(proxy,
			fmt.Sprintf("proxy-%d", proxyPort),
			proxy.Listener,
			handleProxyConn)
	}
	acceptConn(kServer,
		fmt.Sprintf("server-%d", kServer.ServerPort),
		kServer.Listener,
		handleServerConn)
}

func listenPort(port uint16) net.Listener {
	address := fmt.Sprintf("0.0.0.0:%d", port)
	listener, err := net.Listen("tcp", address)
	if err != nil {
		panic(errors.Wrap(err, fmt.Sprintf("Listen on address '%s' failed", address)))
	}
	return listener
}

func acceptConn(obj interface{}, work string, listener net.Listener, handler ConnHandler) {
	failCount := 0
	for {
		conn, err := listener.Accept()
		if err != nil {
			failCount = failCount + 1
			logError(err, fmt.Sprintf("%dth accept connection for '%s' failed", failCount, work))
			continue
		}
		go handler(obj, conn)
	}
}

func handleServerConn(obj interface{}, conn net.Conn) {
	ctx := &ServerConnContext{
		Server:    obj.(*Server),
		Host:      nil,
		Conn:      conn,
		Err:       nil,
		PktHeader: nil,
	}
	defer releaseServerConn(ctx)

	var pktHeaderBytes = make([]byte, PacketHeaderSize)
	for {
		ctx.Err = readFixedData(conn, pktHeaderBytes, getHostReadTimeout(ctx))
		if ctx.Err != nil {
			return
		}

		ctx.PktHeader = (*PacketHeader)(BytesToStruct(pktHeaderBytes))

		if ctx.PktHeader.Magic != PacketMagic ||
			ctx.PktHeader.Version > PacketVersion ||
			int(ctx.PktHeader.ContentSize) > ctx.Server.MaxPacketBodySize {
			ctx.Err = errors.New(fmt.Sprint("Incorrect packet header: ", ctx.PktHeader))
			return
		}

		switch ctx.PktHeader.OperationType {
		case OptTypeHostConn:
			ctx.Err = handleHostConn(ctx)
			break
		case OptTypeHostSyn:
			ctx.Err = handleHostSync(ctx)
			break
		case OptTypeLineConn:
			ctx.Err = handleLineConn(ctx)
			break
		default:
			ctx.Err = errors.New(fmt.Sprint("Incorrect 'OperationType' from packet header: ", ctx.PktHeader))
		}
		if ctx.Err != nil {
			return
		}
	}
}

func releaseServerConn(ctx *ServerConnContext) {
	closeConn(ctx.Conn)
	if ctx.Host != nil && ctx.HostConnected {
		ctx.Host.Connected = 0
		ctx.Host.Conn = nil
	}
	if ctx.Err == nil {
		ctx.Err = recover()
		if ctx.Err != nil {
			debug.PrintStack()
		}
	}
	if ctx.Err == nil {
		ctx.Err = ""
	}
	logError(ctx.Err,
		fmt.Sprintf("Host '%s' conn '%s' released",
			getHostName(ctx),
			getConnInfo(ctx.Conn)))
}

func handleHostConn(ctx *ServerConnContext) error {
	var err error

	contentSize := int(ctx.PktHeader.ContentSize)
	contentBytes := make([]byte, contentSize)

	err = readFixedData(ctx.Conn, contentBytes, getHostReadTimeout(ctx))
	if err != nil {
		return err
	}

	hostConnBody := &HostConnBody{}
	err = json.Unmarshal(contentBytes, hostConnBody)
	if err != nil {
		return errors.Wrap(err,
			fmt.Sprintf("Parse the json '%s' to 'HostConnBody' failed", contentBytes))
	}

	host := ctx.Server.Hosts[hostConnBody.HostName]
	if host == nil {
		return errors.New(
			fmt.Sprintf("Could not found the host with name: '%s'", hostConnBody.HostName))
	}

	ctx.Host = host

	if !hostConnBody.CheckTimestamp(ctx.Server.SignTimeout) {
		return errors.New(
			fmt.Sprintf("Incorrect timestamp '%s'", contentBytes))
	}

	if !hostConnBody.CheckSign(host.SecretKey) {
		return errors.New(
			fmt.Sprintf("Incorrect sign '%s'", contentBytes))
	}

	if !atomic.CompareAndSwapInt32(&host.Connected, 0, 1) {
		return errors.New(
			fmt.Sprintf("Host '%s' was already been connected by '%s'",
				hostConnBody.HostName,
				getConnInfo(host.Conn)))
	}
	host.Conn = ctx.Conn
	ctx.HostConnected = true
	return hostAck(ctx)
}

func handleHostSync(ctx *ServerConnContext) error {
	if ctx.PktHeader.ContentSize != 0 {
		return errors.New(fmt.Sprint("Incorrect 'ContentSize' from packet header:", ctx.PktHeader))
	}
	if ctx.Host.Connected == 0 {
		return errors.New("Unconnected host")
	}
	return hostAck(ctx)
}

func hostAck(ctx *ServerConnContext) error {
	hostAckHeader := &PacketHeader{
		Magic:         PacketMagic,
		Version:       PacketVersion,
		OperationType: OptTypeHostAck,
		ContentSize:   0,
	}
	return writeFixedData(ctx.Conn,
		StructToBytes(unsafe.Pointer(hostAckHeader), PacketHeaderSize),
		getHostWriteTimeout(ctx))
}

func handleLineConn(ctx *ServerConnContext) error {
	var err error

	contentSize := int(ctx.PktHeader.ContentSize)
	contentBytes := make([]byte, contentSize)

	err = readFixedData(ctx.Conn, contentBytes, getHostReadTimeout(ctx))
	if err != nil {
		return err
	}

	lineConnBody := &LineConnBody{}
	err = json.Unmarshal(contentBytes, lineConnBody)
	if err != nil {
		return errors.Wrap(err,
			fmt.Sprintf("Parse the json '%s' to 'LineConnBody' failed", contentBytes))
	}

	host := ctx.Server.Hosts[lineConnBody.HostName]
	if host == nil {
		return errors.New(
			fmt.Sprintf("Could not found the host with name: '%s'", lineConnBody.HostName))
	}

	ctx.Host = host

	if !lineConnBody.CheckTimestamp(ctx.Server.SignTimeout) {
		return errors.New(
			fmt.Sprintf("Incorrect timestamp '%s'", contentBytes))
	}

	if !lineConnBody.CheckSign(host.SecretKey) {
		return errors.New(
			fmt.Sprintf("Incorrect sign '%s'", contentBytes))
	}

	host.LineTreeLock.Lock()
	lineItem := host.Lines.Get(&Line{LineNum: lineConnBody.LineNum})
	host.LineTreeLock.Unlock()
	if lineItem == nil {
		return errors.New(
			fmt.Sprintf("Could not found the line '%d' from host: '%s'",
				lineConnBody.LineNum,
				lineConnBody.HostName))
	}

	line := lineItem.(*Line)
	if !atomic.CompareAndSwapInt32(&line.Connected, 0, 1) {
		return errors.New(
			fmt.Sprintf("Line '%d' of host '%s' was already been connected by '%s'",
				lineConnBody.LineNum,
				lineConnBody.HostName,
				getConnInfo(line.HostConn)))
	}
	line.HostConn = ctx.Conn
	line.ConnChannel <- true
	return transportConnData(
		line.HostConn,
		line.ProxyConn,
		line.Host.ReadTimeout,
		line.Proxy.WriteTimeout,
		line.Server.TransportBufferSize)
}

func handleProxyConn(obj interface{}, conn net.Conn) {
	proxy := obj.(*Proxy)
	host := proxy.Host
	line := &Line{
		Server:      proxy.Server,
		Proxy:       proxy,
		Host:        host,
		LineNum:     atomic.AddInt32(&proxy.Host.LineSequence, 1),
		ProxyConn:   conn,
		HostConn:    nil,
		Connected:   0,
		ConnChannel: make(chan bool),
		Err:         nil,
	}

	defer releaseProxyConn(line)

	host.LineTreeLock.Lock()
	succeed := host.Lines.Insert(line)
	host.LineTreeLock.Unlock()
	if !succeed {
		line.Err = errors.New(fmt.Sprintf(
			"Lines of host '%s' was very busy", host.HostName))
		return
	}

	line.Err = lineCall(line)
	if line.Err != nil {
		return
	}

	succeed = waitLineCallback(line)
	if !succeed {
		line.Err = errors.New(fmt.Sprintf(
			"Wait line '%d' of host '%s' callback failed", line.LineNum, host.HostName))
		return
	}

	line.Err = transportConnData(
		line.ProxyConn,
		line.HostConn,
		line.Proxy.ReadTimeout,
		line.Host.WriteTimeout,
		line.Server.TransportBufferSize)
}

func lineCall(line *Line) error {
	hostConn := line.Host.Conn
	if hostConn == nil {
		return errors.New(fmt.Sprintf("Host '%s' was not connected", line.Host.HostName))
	}

	lineCallBody := &LineCallBody{
		LineNum:       line.LineNum,
		ProxyPort:     line.Proxy.ProxyPort,
		Network:       line.Proxy.Network,
		TargetAddress: line.Proxy.TargetAddress,
	}
	lineCallBody.GenerateSign(line.Host.SecretKey)
	lineCallBodyBytes, err := json.Marshal(lineCallBody)
	if err != nil {
		return err
	}
	lineCallBodySize := len(lineCallBodyBytes)

	lineCallHeader := &PacketHeader{
		Magic:         PacketMagic,
		Version:       PacketVersion,
		OperationType: OptTypeLineCall,
		ContentSize:   int32(lineCallBodySize),
	}
	lineCallHeaderBytes := StructToBytes(unsafe.Pointer(lineCallHeader), PacketHeaderSize)

	lineCallSize := PacketHeaderSize + lineCallBodySize
	lineCallBytes := make([]byte, lineCallSize)
	copy(lineCallBytes, lineCallHeaderBytes)
	copy(lineCallBytes[PacketHeaderSize:], lineCallBodyBytes)

	return writeFixedData(hostConn, lineCallBytes, line.Host.WriteTimeout)
}

func waitLineCallback(line *Line) bool {
	select {
	case <-time.After(line.Host.ReadTimeout):
		return false
	case success := <-line.ConnChannel:
		return success
	}
}

func releaseProxyConn(line *Line) {
	closeConn(line.ProxyConn)
	line.Host.LineTreeLock.Lock()
	line.Host.Lines.DeleteWithKey(line)
	line.Host.LineTreeLock.Unlock()
	if line.Err == nil {
		line.Err = recover()
		if line.Err != nil {
			debug.PrintStack()
		}
	}
	if line.Err == nil {
		line.Err = ""
	}
	logError(line.Err,
		fmt.Sprintf("Line: '%d', route: '%s %d->%s %s', proxy conn: '%s', host conn: '%s' released",
			line.LineNum,
			line.Proxy.Network,
			line.Proxy.ProxyPort,
			line.Host.HostName,
			line.Proxy.TargetAddress,
			getConnInfo(line.ProxyConn),
			getConnInfo(line.HostConn)))
}

func transportConnData(
	readConn net.Conn,
	writeConn net.Conn,
	readTimeout time.Duration,
	writeTimeout time.Duration,
	bufferSize int) error {
	buffer := make([]byte, bufferSize)
	for {
		err := readConn.SetReadDeadline(time.Now().Add(readTimeout))
		if err != nil {
			return err
		}
		readSize, err := readConn.Read(buffer)
		if err != nil {
			return err
		}
		err = writeFixedData(writeConn, buffer[0:readSize], writeTimeout)
		if err != nil {
			return err
		}
	}
}

func getHostReadTimeout(ctx *ServerConnContext) time.Duration {
	if ctx.Host == nil {
		return ctx.Server.HostReadTimeout
	}
	return ctx.Host.ReadTimeout
}

func getHostWriteTimeout(ctx *ServerConnContext) time.Duration {
	if ctx.Host == nil {
		return ctx.Server.HostWriteTimeout
	}
	return ctx.Host.WriteTimeout
}

func getHostName(ctx *ServerConnContext) string {
	if ctx.Host == nil {
		return "__unknown__"
	}
	return ctx.Host.HostName
}

func readFixedData(conn net.Conn, bytes []byte, timeout time.Duration) error {
	if bytes == nil {
		return nil
	}
	size := len(bytes)
	if size == 0 {
		return nil
	}
	readSize := 0
	for readSize < size {
		err := conn.SetReadDeadline(time.Now().Add(timeout))
		if err != nil {
			return err
		}
		n, err := conn.Read(bytes[readSize:size])
		if err != nil {
			return err
		}
		readSize += n
	}
	return nil
}

func writeFixedData(conn net.Conn, bytes []byte, timeout time.Duration) error {
	if bytes == nil {
		return nil
	}
	size := len(bytes)
	if size == 0 {
		return nil
	}
	writeSize := 0
	for writeSize < size {
		err := conn.SetWriteDeadline(time.Now().Add(timeout))
		if err != nil {
			return err
		}
		n, err := conn.Write(bytes[writeSize:size])
		if err != nil {
			return err
		}
		writeSize += n
	}
	return nil
}

func getConnInfo(conn net.Conn) string {
	if conn == nil {
		return "__unconnected__"
	}
	return fmt.Sprintf("%s %s->%s %s",
		conn.LocalAddr().Network(),
		conn.LocalAddr().String(),
		conn.RemoteAddr().Network(),
		conn.RemoteAddr().String())
}

func parseDurationConfig(name string, value string) time.Duration {
	duration, err := time.ParseDuration(value)
	if err != nil {
		panic(errors.Wrap(err, fmt.Sprintf("Incorrect duration property '%s': %s", name, value)))
	}
	return duration
}

func parseObjDurationConfig(objName string, name string, value string, defaultValue time.Duration) time.Duration {
	if strings.Compare(value, "") == 0 {
		return defaultValue
	}
	duration, err := time.ParseDuration(value)
	if err != nil {
		panic(errors.Wrap(err, fmt.Sprintf("Incorrect duration property '%s.%s': %s", objName, name, value)))
	}
	return duration
}

func closeFile(file *os.File) {
	err := file.Close()
	if err != nil {
		logError(err, fmt.Sprintf("Close file '%s' failed", file.Name()))
	}
}

func closeConn(conn net.Conn) {
	if conn == nil {
		return
	}
	err := conn.Close()
	if err != nil {
		logError(err, fmt.Sprintf("Close conn failed"))
	}
}

func lineCompare(a, b rbtree.Item) int {
	lineNumA := a.(*Line).LineNum
	lineNumB := b.(*Line).LineNum
	if lineNumA > lineNumB {
		return 1
	}
	if lineNumA < lineNumB {
		return -1
	}
	return 0
}

func logError(err interface{}, message string) {
	_, _ = fmt.Fprintln(
		os.Stderr,
		fmt.Sprintf("[%s] %s:",
			time.Now().Format("2006-01-02T03:04:05.999999999Z07:00"),
			message),
		err)
}

func main() {
	initConfig()
	listen()
}
